package com.uinnova.product.eam.base.web.http.response;

import lombok.AllArgsConstructor;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;

/**
 * 自定义response 主要重写了写入方法，使其http输出流保持的同时新增一个字节输出流保证程序可以读取response种的responsebody
 * 中心逻辑就是将写入到ServletOutputStream修改为写入到ByteArrayOutputStream；
 * 准备flush时再从ByteArrayOutputStream拷贝到ServletOutputStream
 * 
 * @author zhaoyujie
 *
 */
public class CustomHttpServletResponseWrapper extends HttpServletResponseWrapper {
	private ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
	private HttpServletResponse response;

	public CustomHttpServletResponseWrapper(HttpServletResponse response) {
		super(response);
		this.response = response;
	}

	/**
	 * 获取请求体输出字节流
	 * 
	 * @return
	 */
	public byte[] getBody() {
		return byteArrayOutputStream.toByteArray();
	}

	@Override
	public ServletOutputStream getOutputStream() {
		return new ServletOutputStreamWrapper(this.byteArrayOutputStream, this.response);
	}

	@Override
	public PrintWriter getWriter() throws IOException {
		return new PrintWriter(new OutputStreamWriter(getOutputStream(), this.response.getCharacterEncoding()));
	}

	/**
	 * 
	 * @author zhaoyujie
	 *
	 */
	@AllArgsConstructor
	private static class ServletOutputStreamWrapper extends ServletOutputStream {

		private ByteArrayOutputStream outputStream;
		private HttpServletResponse response;

		@Override
		public boolean isReady() {
			return true;
		}

		@Override
		public void setWriteListener(WriteListener listener) {

		}

		@Override
		public void write(int b) throws IOException {
			this.outputStream.write(b);
		}

		/**
		 * 保障复写前逻辑，最后要写到ServletOutputStream，否则无法传递client
		 */
		@Override
		public void flush() throws IOException {
			if (!this.response.isCommitted()) {
				byte[] body = this.outputStream.toByteArray();
				ServletOutputStream outputStream = this.response.getOutputStream();
				outputStream.write(body);
				outputStream.flush();
			}
		}
	}
}
